#include <iostream>
#include <algorithm>
#include <chrono>
#include <stdio.h>
#include <cuda_runtime.h>
#include <cstring>

// 定义常量
// #define DATASIZE 11
// #define BLOCK_SIZE 3
// #define MAXBIT 1
#define MAXBIT 32
#define DATASIZE 1000000
#define BLOCK_SIZE 256


using namespace std::chrono;

bool validate(unsigned int *a, unsigned int* b, int n){
    for (int i = 0; i < n; i++) {
        if(a[i] != b[i]){
            printf("index = %d is not matched.\n", i);
            return false;
        }
    }
    return true;
}

void print_array(unsigned int* arr, int n, const char* label) {
    if(n > 20) return;
    printf("\n%s\n", label);
    for (int i = 0; i < n; i++) {
        printf("%d  ", arr[i]);
    }
    printf("\n");
}

void cpu_radix_sort_with_index(unsigned int* arr, unsigned int* arr_pos,
                               unsigned int* sorted_arr, unsigned int* original_indices,
                               int n) {
    unsigned int* values_curr = new unsigned int[n];
    unsigned int* values_next = new unsigned int[n];
    unsigned int* pos_curr = new unsigned int[n];
    unsigned int* pos_next = new unsigned int[n];
    std::copy(arr, arr + n, values_curr);
    std::copy(arr_pos, arr_pos + n, pos_curr);
    for (int bit = 0; bit < MAXBIT; bit++) {
        int count[2] = {0, 0};
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            count[digit]++;
        }
        int offset[2] = {0, count[0]};
        for (int i = 0; i < n; i++) {
            int digit = (values_curr[i] >> bit) & 1;
            int dest_idx = offset[digit]++;
            values_next[dest_idx] = values_curr[i];
            pos_next[dest_idx] = pos_curr[i];
        }
        std::swap(values_curr, values_next);
        std::swap(pos_curr, pos_next);
    }
    // 循环结束后，最终结果在 values_curr 和 pos_curr 中
    // 因为每次 swap 后，curr 总是指向最新结果
    std::copy(values_curr, values_curr + n, sorted_arr);
    std::copy(pos_curr, pos_curr + n, original_indices);

    delete[] values_curr;
    delete[] values_next;
    delete[] pos_curr;
    delete[] pos_next;
}

void computerBit(unsigned int* d_in, unsigned int* d_in_bit, int bit, int n){
    for(int i = 0; i < n ; ++i){
         d_in_bit[i] = (d_in[i] >> bit) & 1;
    }
}

// 提取当前 bit
__global__ void computeBitValues(unsigned int* d_in, unsigned int* d_in_bit, int bit, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        d_in_bit[idx] = (d_in[idx] >> bit) & 1;
    }
}

// 块内前缀和 + 输出块总和
__global__
void scanBlockSums(const unsigned int* d_in, unsigned int* d_out, unsigned int* blockSums, int n) {
    extern __shared__ unsigned int shared_data[];
    int tid = threadIdx.x;
    int globalId = blockDim.x * blockIdx.x + threadIdx.x;
    bool isValid = globalId < n;

    // 初始化共享内存
    shared_data[tid] = isValid ? d_in[globalId] : 0;
    __syncthreads();

    // Hillis-Steele inclusive scan
    int pin = 0, pout = 1;
    for (int s = 1; s < blockDim.x; s <<= 1) {
        if (tid < s) {
            shared_data[pout * blockDim.x + tid] = shared_data[pin * blockDim.x + tid];
        }
        if (tid + s < blockDim.x) {
            shared_data[pout * blockDim.x + tid + s] = 
                shared_data[pin * blockDim.x + tid] + shared_data[pin * blockDim.x + tid + s];
        }
        __syncthreads();
        pin = 1 - pin; pout = 1 - pout;
    }

    // 写出结果
    if (isValid) {
        d_out[globalId] = shared_data[pin * blockDim.x + tid];
    }

    // 只有当该 block 有有效数据时才写 blockSum
    if (tid == blockDim.x - 1) {
        // 计算该 block 的实际 sum（只包含有效元素）
        unsigned int sum = 0;
        int start = blockDim.x * blockIdx.x;
        int end = min(start + blockDim.x, n);
        for (int i = start; i < end; ++i) {
            sum += d_in[i];
        }
        blockSums[blockIdx.x] = sum;
    }
}

// 反转数组
__global__ void reverseArray(unsigned int* input, unsigned int* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        output[n - 1 - idx] = input[idx];
    }
}

void rearrange(unsigned int* d_in, unsigned int* d_in_pos, unsigned int* d_in_bit, 
                                unsigned int* d_prefix_sum, unsigned int* d_suffix_sum, 
                                unsigned int* d_tmp, unsigned int* d_tmp_pos, int n){
    for(int i = 0; i < n; ++i){
       if (d_in_bit[i] == 0) {
            d_tmp[d_prefix_sum[i] - 1] = d_in[i];
            d_tmp_pos[d_prefix_sum[i] - 1] = d_in_pos[i];
        } else {
            d_tmp[n - d_suffix_sum[i]] = d_in[i];
            d_tmp_pos[n - d_suffix_sum[i]] = d_in_pos[i];
        } 
    }                                 
}

// 重排数组
__global__ void rearrangeArray(unsigned int* d_in, unsigned int* d_in_pos, unsigned int* d_in_bit, 
                                unsigned int* d_prefix_sum, unsigned int* d_suffix_sum, 
                                unsigned int* d_tmp, unsigned int* d_tmp_pos, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        if (d_in_bit[idx] == 0) {
            d_tmp[d_prefix_sum[idx] - 1] = d_in[idx];
            d_tmp_pos[d_prefix_sum[idx] - 1] = d_in_pos[idx];
        } else {
            d_tmp[n - d_suffix_sum[idx]] = d_in[idx];
            d_tmp_pos[n - d_suffix_sum[idx]] = d_in_pos[idx];
        }
    }
}

// CPU 前缀和，用于验证，逻辑是统计 0 的数量
void prefixSum(unsigned int* d_input, unsigned int* d_output, int n){
    int tmp = 0;
    for(int i = 0; i < n ; ++i){
        tmp = (d_input[i] == 0) ? 1 : 0; // 统计 0
        if(i == 0){
            d_output[i] = tmp;
        }else{
            d_output[i] = tmp + d_output[i - 1];
        }
    }
}

// CPU 后缀和，用于验证，逻辑是统计 1 的数量
void suffixSum(unsigned int* d_input, unsigned int* d_output, int n){
    for(int i = n - 1; i >= 0; --i){
        if(i == (n-1)){
            d_output[i] = (d_input[i] == 1) ? 1 : 0; // 统计 1
        }else{
            d_output[i] = ((d_input[i] == 1) ? 1 : 0) + d_output[i + 1];
        }
    }
}

// 辅助kernel：对输入取反 (bit=0 -> 1, bit=1 -> 0)
__global__ void kernel_inverse(unsigned int* input, unsigned int* output, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx < n) {
        output[idx] = (input[idx] == 0) ? 1 : 0;
    }
}

// 新增：并行加偏移 kernel（所有 block 同时加）
__global__ void addOffsetParallel(unsigned int* output, unsigned int* blockOffsets, int n) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    if (idx >= n) return;
    int blockId = blockIdx.x;
    if (blockId == 0) return;
    output[idx] += blockOffsets[blockId];
}

__global__ void shiftRightExclusive(unsigned int* array, int n) {
    if (n <= 1) return;
    // 从最后一个元素开始，依次向前移动
    for (int i = n - 1; i > 0; i--) {
        array[i] = array[i - 1];
    }
    // 第一个元素设为0
    array[0] = 0;
}

// 辅助函数：CPU 版本的 kernel_inverse
// 将输入数组中 0 变为 1，1 变为 0
void cpu_kernel_inverse(unsigned int* input, unsigned int* output, int n) {
    for (int i = 0; i < n; ++i) {
        output[i] = (input[i] == 0) ? 1 : 0;
    }
}

// 辅助函数：CPU 版本的 shiftRightExclusive
// 将数组右移一位，第一个元素设为0，实现前缀和的“排他性”
void cpu_shiftRightExclusive(unsigned int* array, int n) {
    if (n <= 1) return;
    // 从后往前移动，避免覆盖
    for (int i = n - 1; i > 0; --i) {
        array[i] = array[i - 1];
    }
    array[0] = 0;
}

// 辅助函数：CPU 版本的 addOffsetParallel
// 为每个元素加上其所在块的偏移量
void cpu_addOffsetParallel(unsigned int* output, unsigned int* blockOffsets, int n, int blockSize) {
    for (int i = 0; i < n; ++i) {
        int blockId = i / blockSize;
        if (blockId > 0) {
            output[i] += blockOffsets[blockId];
        }
    }
}

void cpu_scanBlockSums(unsigned int* input, unsigned int* output, unsigned int* blocksum, int n, int blockSize) {
    for (int blockId = 0; blockId * blockSize < n; ++blockId) {
        int startIdx = blockId * blockSize;
        int endIdx = std::min((blockId + 1) * blockSize, n);
        unsigned int sum = 0;
        for (int i = startIdx; i < endIdx; ++i) {
            sum += input[i];
            output[i] = sum;
        }
        blocksum[blockId] = sum;
    }
}

void cpu_scanBlockSumsNormal(unsigned int* input, unsigned int* output, int n) {
    output[0] = input[0];
    for (int i = 1; i < n - 1; ++i) {
        output[i] = input[i] + output[i - 1];
    }
}

// 完整的、可扩展的 multi-block inclusive prefix sum
void inclusivePrefixSum(unsigned int* d_input, unsigned int* d_output, int n, bool inverse) {
    int numBlocks = (n + BLOCK_SIZE - 1) / BLOCK_SIZE;
    unsigned int *d_temp, *d_blockSums, *d_blockOffsets;

    // 仅当 n > 0 时分配内存
    if (n > 0) {
        cudaMalloc(&d_temp,        n * sizeof(unsigned int));
        cudaMalloc(&d_blockSums,   numBlocks * sizeof(unsigned int));
        cudaMalloc(&d_blockOffsets, numBlocks * sizeof(unsigned int));
    }
    dim3 block(BLOCK_SIZE);
    dim3 grid(numBlocks);

    // Step 1: 如果需要取反（统计 bit=0），先处理
    if (n > 0 && inverse) {
        kernel_inverse<<<grid, block>>>(d_input, d_temp, n);
        cudaDeviceSynchronize();

        // to_do 在此调用cpu版本的kernel_inverse算法，然后与kernel_inverse产生的结果对比。
        // unsigned int *h_input = new unsigned int[n];
        // unsigned int *h_temp_gpu = new unsigned int[n];
        // unsigned int *h_temp_cpu = new unsigned int[n];
        // cudaMemcpy(h_input, d_input, n * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // cudaMemcpy(h_temp_gpu, d_temp, n * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // cpu_kernel_inverse(h_input, h_temp_cpu, n);
        // print_array(h_temp_cpu, n, "step1 cpu:");
        // print_array(h_temp_gpu, n, "step1 gpu:");
        // printf("kernel_inverse : %s\n", validate(h_temp_gpu, h_temp_cpu, n) ? "CORRECT" : "INCORRECT");
        // delete[] h_input; delete[] h_temp_gpu; delete[] h_temp_cpu;
    }else if (n > 0) {
        cudaMemcpy(d_temp, d_input, n * sizeof(unsigned int), cudaMemcpyDeviceToDevice);
    } 

    // Step 2: 块内 inclusive 扫描，输出 blockSums
    if (n > 0) {
        scanBlockSums<<<grid, block, 2 * BLOCK_SIZE * sizeof(unsigned int)>>>(d_temp, d_output, d_blockSums, n);
        cudaDeviceSynchronize();
    }

    // to_do 在此调用cpu版本的scanBlockSums算法，然后与scanBlockSums产生的结果对比。
    // unsigned int *h_temp = new unsigned int[n];
    // unsigned int *h_output_gpu = new unsigned int[n];
    // unsigned int *h_output_cpu = new unsigned int[n];
    // unsigned int *h_blockSums_gpu = new unsigned int[numBlocks];
    // unsigned int *h_blockSums_cpu = new unsigned int[numBlocks];
    // cudaMemcpy(h_temp, d_temp, n * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cudaMemcpy(h_output_gpu, d_output, n * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cudaMemcpy(h_blockSums_gpu, d_blockSums, numBlocks * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cpu_scanBlockSums(h_temp, h_output_cpu, h_blockSums_cpu, n, BLOCK_SIZE);
    // print_array(h_output_cpu, n, "step2 cpu:");
    // print_array(h_output_gpu, n, "step2 gpu:");
    // printf("scanBlockSums (output): %s\n", validate(h_output_gpu, h_output_cpu, n) ? "CORRECT" : "INCORRECT");
    // print_array(h_blockSums_cpu, numBlocks, "step2 blocksum cpu:");
    // print_array(h_blockSums_gpu, numBlocks, "step2 blocksum gpu:");
    // printf("scanBlockSums (blockSums): %s\n", validate(h_blockSums_gpu, h_blockSums_cpu, numBlocks - 1)  ? "CORRECT" : "INCORRECT");
    // delete[] h_temp; delete[] h_output_gpu; delete[] h_blockSums_gpu;

    // Step 3: 如果只有一个 block，无需后续操作
    if (numBlocks <= 1 || n == 0) {
        if (d_temp) cudaFree(d_temp);
        if (d_blockSums) cudaFree(d_blockSums);
        if (d_blockOffsets) cudaFree(d_blockOffsets);
        return;
    }

    // Step 4: 递归处理 blockSums
    // 注意：这里对 d_blockSums 递归调用 inclusivePrefixSum，结果存入 d_blockOffsets
    inclusivePrefixSum(d_blockSums, d_blockOffsets, numBlocks, false);
    cudaDeviceSynchronize();

    // to_do 在此调用cpu版本的scanBlockSums算法，然后与scanBlockSums产生的结果对比。
    // unsigned int *h_blockOffsets_gpu = new unsigned int[numBlocks];
    // unsigned int *h_blockOffsets_cpu = new unsigned int[numBlocks];
    // cudaMemcpy(h_blockOffsets_gpu, d_blockOffsets, numBlocks * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cpu_scanBlockSumsNormal(h_blockSums_cpu, h_blockOffsets_cpu, numBlocks);
    // print_array(h_blockOffsets_cpu, numBlocks, "step4 scanblocksum cpu:");
    // print_array(h_blockOffsets_gpu, numBlocks, "step4 scanblocksum gpu:");
    // printf("scanBlockSums (blockOffsets): %s\n", validate(h_blockOffsets_gpu, h_blockOffsets_cpu, numBlocks - 1) ? "CORRECT" : "INCORRECT");
    // delete[] h_blockSums_cpu;delete[] h_blockOffsets_gpu;

    // Step 5: 将 inclusive 的 blockOffsets 变为 exclusive
    // d_blockOffsets 现在是 inclusive 的，例如 [3, 4, 7] 将其变为 exclusive，即 [0, 3, 4]
    shiftRightExclusive<<<1, 1>>>(d_blockOffsets, numBlocks);
    cudaDeviceSynchronize();

    // to_do 在此调用cpu版本的shiftRightExclusive算法，然后与shiftRightExclusive产生的结果对比。
    // unsigned int *h_blockOffsets_excl_gpu = new unsigned int[numBlocks];
    // cudaMemcpy(h_blockOffsets_excl_gpu, d_blockOffsets, numBlocks * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cpu_shiftRightExclusive(h_blockOffsets_cpu, numBlocks);
    // print_array(h_blockOffsets_cpu, numBlocks, "step5 shift cpu:");
    // print_array(h_blockOffsets_excl_gpu, numBlocks, "step5 shift gpu:");
    // printf("shiftRightExclusive: %s\n", validate(h_blockOffsets_excl_gpu, h_blockOffsets_cpu, numBlocks) ? "CORRECT" : "INCORRECT");
    // delete[] h_blockOffsets_excl_gpu;

    // Step 6: 并行地为每个 block 加上 offset
    addOffsetParallel<<<grid, block>>>(d_output, d_blockOffsets, n);
    cudaDeviceSynchronize();

    // to_do 在此调用cpu版本的addOffsetParallel算法(需要新补充此方法)，然后与addOffsetParallel产生的结果对比。
    // unsigned int *h_output_final_gpu = new unsigned int[n];
    // cudaMemcpy(h_output_final_gpu, d_output, n * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    // cpu_addOffsetParallel(h_output_cpu, h_blockOffsets_cpu, n, BLOCK_SIZE);
    // print_array(h_output_cpu, n, "step6 addoffset cpu:");
    // print_array(h_output_final_gpu, n, "step6 addoffset gpu:");
    // printf("addOffsetParallel: %s\n", validate(h_output_final_gpu, h_output_cpu, n) ? "CORRECT" : "INCORRECT");
    // delete[] h_output_final_gpu; delete[] h_blockOffsets_cpu;delete[] h_output_cpu;

    cudaFree(d_temp);
    cudaFree(d_blockSums);
    cudaFree(d_blockOffsets);
}

// 后缀和
void inclusiveSuffixSum(unsigned int* d_input, unsigned int* d_output, int n, bool inverse) {
    unsigned int *d_rev_input = nullptr, *d_rev_sum = nullptr;
    if (n == 0) {
        return;
    }
    cudaMalloc(&d_rev_input, n * sizeof(unsigned int));
    cudaMalloc(&d_rev_sum,    n * sizeof(unsigned int));
    dim3 block(BLOCK_SIZE);
    dim3 grid((n + BLOCK_SIZE - 1) / BLOCK_SIZE);

    reverseArray<<<grid, block>>>(d_input, d_rev_input, n);
    cudaDeviceSynchronize();

    inclusivePrefixSum(d_rev_input, d_rev_sum, n, false);
    cudaDeviceSynchronize();

    reverseArray<<<grid, block>>>(d_rev_sum, d_output, n);
    cudaDeviceSynchronize();

    cudaFree(d_rev_input);
    cudaFree(d_rev_sum);
}


void gpu_radix_sort(unsigned int* h_in, unsigned int* h_in_pos, 
                    unsigned int* h_out, unsigned int* h_out_pos, int N) {
    unsigned int *d_in, *d_in_pos, *d_tmp, *d_tmp_pos, *d_in_bit, *d_prefix_sum, *d_suffix_sum;
    cudaMalloc(&d_in, N * sizeof(unsigned int));
    cudaMalloc(&d_in_pos, N * sizeof(unsigned int));
    cudaMalloc(&d_tmp, N * sizeof(unsigned int));
    cudaMalloc(&d_tmp_pos, N * sizeof(unsigned int));
    cudaMalloc(&d_in_bit, N * sizeof(unsigned int));
    cudaMalloc(&d_prefix_sum, N * sizeof(unsigned int));
    cudaMalloc(&d_suffix_sum, N * sizeof(unsigned int));
    cudaMemcpy(d_in, h_in, N * sizeof(unsigned int), cudaMemcpyHostToDevice);
    cudaMemcpy(d_in_pos, h_in_pos, N * sizeof(unsigned int), cudaMemcpyHostToDevice);
    int numBlocks = (N + BLOCK_SIZE - 1) / BLOCK_SIZE;
    for (int bit = 0; bit < MAXBIT; bit++) {
        // 计算每一位的值 0或1
        computeBitValues<<<numBlocks, BLOCK_SIZE>>>(d_in, d_in_bit, bit, N);
        cudaDeviceSynchronize();

        // unsigned int *cpu_in = new unsigned int[N];
        // unsigned int *h_in_bit = new unsigned int[N];
        // unsigned int *cpu_in_bit = new unsigned int[N];
        // cudaMemcpy(cpu_in, d_in, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // computerBit(cpu_in, cpu_in_bit, bit, N);
        // cudaMemcpy(h_in_bit, d_in_bit, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // print_array(cpu_in_bit, N, "cpu bit:");
        // print_array(h_in_bit, N, "gpu bit:");
        // printf("computeBit : %s\n", validate(h_in_bit, cpu_in_bit, N)? "CORRECT" : "INCORRECT");

        // 统计前缀和，所以 inverse=true
        inclusivePrefixSum(d_in_bit, d_prefix_sum, N, true);
        cudaDeviceSynchronize();

        // unsigned int *h_prefix = new unsigned int[N];
        // unsigned int *cpu_prefix = new unsigned int[N];
        // prefixSum(cpu_in_bit, cpu_prefix, N);
        // cudaMemcpy(h_prefix, d_prefix_sum, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // print_array(cpu_prefix, N, "cpu prefix:");
        // print_array(h_prefix, N, "gpu prefix:");
        // printf("prefixSum : %s\n", validate(h_prefix, cpu_prefix, N)? "CORRECT" : "INCORRECT");

        // 统计后缀和，inverse=false
        inclusiveSuffixSum(d_in_bit, d_suffix_sum, N, false);
        cudaDeviceSynchronize();

        // unsigned int *h_suffix = new unsigned int[N];
        // unsigned int *cpu_suffix = new unsigned int[N];
        // suffixSum(cpu_in_bit, cpu_suffix, N);
        // cudaMemcpy(h_suffix, d_suffix_sum, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // print_array(cpu_suffix, N, "cpu suffix:");
        // print_array(h_suffix, N, "gpu suffix:");
        // printf("suffixSum : %s\n", validate(h_suffix, cpu_suffix, N)? "CORRECT" : "INCORRECT");
        // 根据前缀和和后缀和重新排列数组
        rearrangeArray<<<numBlocks, BLOCK_SIZE>>>(d_in, d_in_pos, d_in_bit, d_prefix_sum, d_suffix_sum, d_tmp, d_tmp_pos, N);
        cudaDeviceSynchronize();

        // unsigned int *h_rearrange = new unsigned int[N];
        // unsigned int *h_rearrange_pos = new unsigned int[N];
        // unsigned int *cpu_rearrange = new unsigned int[N];
        // unsigned int *cpu_rearrange_pos = new unsigned int[N];
        // unsigned int *cpu_in_pos = new unsigned int[N];
        // cudaMemcpy(cpu_in_pos, d_in_pos, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // rearrange(cpu_in, cpu_in_pos, h_in_bit, cpu_prefix, cpu_suffix, cpu_rearrange, cpu_rearrange_pos, N);
        // cudaMemcpy(h_rearrange, d_tmp, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // cudaMemcpy(h_rearrange_pos, d_tmp_pos, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
        // print_array(cpu_rearrange, N, "cpu rearrange:");
        // print_array(h_rearrange, N, "gpu rearrange:");
        // print_array(cpu_rearrange_pos, N, "cpu rearrange pos:");
        // print_array(h_rearrange_pos, N, "gpu rearrange pos:");
        // printf("rearrange : %s\n", validate(h_rearrange, cpu_rearrange, N)? "CORRECT" : "INCORRECT");
        // printf("rearrange pos : %s\n", validate(h_rearrange_pos, cpu_rearrange_pos, N)? "CORRECT" : "INCORRECT");

        std::swap(d_in, d_tmp);
        std::swap(d_in_pos, d_tmp_pos);
        
        // delete[] cpu_in;delete[] cpu_in_pos;delete[] h_in_bit;delete[] cpu_in_bit; delete[] h_prefix; delete[] cpu_prefix;
        // delete[] h_suffix; delete[] cpu_suffix; delete[] h_rearrange; delete[] h_rearrange_pos; 
        // delete[] cpu_rearrange;delete[] cpu_rearrange_pos;
    }

    cudaMemcpy(h_out, d_in, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);
    cudaMemcpy(h_out_pos, d_in_pos, N * sizeof(unsigned int), cudaMemcpyDeviceToHost);

    cudaFree(d_in);
    cudaFree(d_in_pos);
    cudaFree(d_tmp);
    cudaFree(d_tmp_pos);
    cudaFree(d_in_bit);
    cudaFree(d_prefix_sum);
    cudaFree(d_suffix_sum);
}

int main() {
    // unsigned int* h_in = new unsigned int[DATASIZE] {5, 3, 7, 2, 9}; //, 2 , 8, 4, 6, 0, 3};
    // unsigned int* h_in_pos = new unsigned int[DATASIZE] {0, 1, 2, 3, 4};//, 5 , 6, 7, 8, 9, 10};
    unsigned int* h_in = new unsigned int[DATASIZE];
    unsigned int* h_in_pos = new unsigned int[DATASIZE];
    unsigned int* h_out = new unsigned int[DATASIZE];
    unsigned int* h_out_pos = new unsigned int[DATASIZE];
    unsigned int* cpu_out = new unsigned int[DATASIZE];
    unsigned int* cpu_out_pos = new unsigned int[DATASIZE];

    srand(100);
    for (int i = 0; i < DATASIZE; ++i) {
        h_in[i] = rand() % (DATASIZE + 1);
        h_in_pos[i] = i;
    }    

    auto start_cpu = high_resolution_clock::now();
    cpu_radix_sort_with_index(h_in, h_in_pos, cpu_out, cpu_out_pos, DATASIZE);
    auto end_cpu = high_resolution_clock::now();
    auto cpu_time = duration_cast<milliseconds>(end_cpu - start_cpu).count();

    auto start = high_resolution_clock::now();
    gpu_radix_sort(h_in, h_in_pos, h_out, h_out_pos, DATASIZE);
    auto end = high_resolution_clock::now();
    auto gpu_time = duration_cast<milliseconds>(end - start).count();

    print_array(h_in, DATASIZE, "originally value:");
    print_array(cpu_out, DATASIZE, "cpu sorted value:");
    print_array(h_out, DATASIZE, "gpu sorted value:");
    print_array(cpu_out_pos, DATASIZE, "cpu sorted index:");
    print_array(h_out_pos, DATASIZE, "gpu sorted index:");
    printf("Sort value: %s\n", validate(h_out, cpu_out, DATASIZE) ? "CORRECT" : "INCORRECT");
    printf("Sort index: %s\n", validate(h_out_pos, cpu_out_pos, DATASIZE) ? "CORRECT" : "INCORRECT");
    std::cout << "Data size: " << DATASIZE << " elements" << std::endl;
    std::cout << "GPU  time: " << gpu_time << " ms" << std::endl;
    std::cout << "CPU  time: " << cpu_time << " ms" << std::endl;

    delete[] h_in;
    delete[] h_in_pos;
    delete[] h_out;
    delete[] h_out_pos;
    delete[] cpu_out;
    delete[] cpu_out_pos;
    return 0;
}